import datetime

import numpy as np
import pandas as pd
import joblib
import os

from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, r2_score

from ..features.feature_engineer import FeatureEngineer


class ModelTrainer:
    """
    模型训练器 - 负责训练成绩预测模型
    """

    def __init__(self):
        self.model = None
        self.feature_engineer = FeatureEngineer()


    def prepare_training_data(self):
        """
        准备训练数据
        使用真实的FeatureEngineer计算特征
        """
        print("📊 准备训练数据...")

        # 🎯 第一阶段：先用模拟数据让流程跑通
        training_data = self._generate_sample_data_with_real_features()

        # 🎯 第二阶段：从数据库加载真实数据（注释掉，后续启用）
        # training_data = self._load_real_data_from_db()

        return training_data


    def _generate_sample_data_with_real_features(self):
        """
        生成使用真实特征工程的模拟数据
        """
        np.random.seed(42)
        num_samples = 1000

        data = []
        for i in range(num_samples):
            # 模拟学生数据
            exam_count = np.random.randint(1, 10)
            scores = []
            base_score = np.random.normal(75, 15)  # 基础水平

            for j in range(exam_count):
                # 模拟成绩波动和进步趋势
                score = base_score + np.random.normal(0, 8) + j * 0.8  # 进步趋势
                score = max(0, min(150, score))  # 限制在0-150分
                scores.append(score)

            # 使用真实的 FeatureEngineer 计算特征！
            student_data = {
                'historical_scores': scores,
                'subject_full_score': 150  # 后续需要根据 score中 model层的总分进行调整
            }

            features = self.feature_engineer.extract_features(student_data)

            # 模拟目标分数（基于特征的真实关系）
            next_score = self._simulate_next_score(features, scores)

            data.append({                 # 现在特征数据只是分数，先试模型是否合适
                'features': features,
                'target': next_score,     # 这个标签不太对，怎么会是下一次的成绩呢。以后需要改
                'student_id': f'student_{i}',
                'historical_scores': scores
            })

        print(f"✅ 生成 {len(data)} 条使用真实特征的数据")
        return data

    def _simulate_next_score(self, features, historical_scores):
        """基于特征模拟下一次考试分数"""
        # 使用特征来模拟真实的分数关系
        base_prediction = (
                features['recent_score'] * 0.5 +
                features['avg_score'] * 0.3 +
                features['score_trend'] * 5 +  # 趋势影响
                np.random.normal(0, 4)  # 随机波动
        )

        # 确保分数在合理范围内
        next_score = max(0, min(150, base_prediction))
        return next_score

    def _get_expected_features(self):
        """获取期望的特征列表"""
        feature_categories = self.feature_engineer.get_feature_categories()
        all_features = []
        for category_features in feature_categories.values():
            all_features.extend(category_features)

        # 移除重复项并排序
        expected_features = list(set(all_features))
        expected_features.sort()
        return expected_features

    def _features_to_vector(self, features):
        """将特征字典转换为特征向量"""
        expected_features = self._get_expected_features()

        vector = []
        for feature_name in expected_features:
            vector.append(features.get(feature_name, 0.0))

        return np.array(vector)

    def train(self):
        """
        训练模型
        """
        print("🚀 开始训练模型...")

        try:
            # 1. 准备数据
            training_data = self.prepare_training_data()

            # 2. 提取特征和标签
            X = []
            y = []

            for sample in training_data:
                feature_vector = self._features_to_vector(sample['features'])
                X.append(feature_vector)
                y.append(sample['target'])

            X = np.array(X)
            y = np.array(y)

            print(f"📈 数据形状: X{X.shape}, y{y.shape}")

            # 3. 分割训练集和测试集
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=0.2, random_state=42
            )

            # 4. 训练模型
            print("🤖 训练随机森林模型...")
            self.model = RandomForestRegressor(
                n_estimators=100,
                random_state=42,
                max_depth=10,
                min_samples_split=5
            )

            self.model.fit(X_train, y_train)

            # 5. 评估模型
            train_score = self.model.score(X_train, y_train)
            test_score = self.model.score(X_test, y_test)

            y_pred = self.model.predict(X_test)
            mae = mean_absolute_error(y_test, y_pred)  # 计算平均预测误差

            # 特征重要性
            feature_importance = dict(zip(self._get_expected_features(), self.model.feature_importances_))

            print(f"✅ 训练完成!")
            print(f"   - 训练集 R²: {train_score:.4f}")
            print(f"   - 测试集 R²: {test_score:.4f}")
            print(f"   - 平均绝对误差: {mae:.2f} 分")
            print(f"   - 最重要的特征: {sorted(feature_importance.items(), key=lambda x: x[1], reverse=True)[:3]}")

            # 6. 保存模型
            self.save_model()

            return {
                'success': True,
                'train_score': float(train_score),
                'test_score': float(test_score),
                'mae': float(mae),
                'model_type': 'RandomForestRegressor',
                'feature_importance': feature_importance,
                'training_samples': len(X),
                'feature_count': X.shape[1]
            }

        except Exception as e:
            print(f"❌ 训练失败: {str(e)}")
            return {'success': False, 'error': str(e)}


    def save_model(self, model_path=None):
        """保存训练好的模型"""

        path = model_path or 'ml/models/score_predictor.pkl'

        # 确保目录存在
        os.makedirs(os.path.dirname(path), exist_ok=True)

        try:
            # 保存模型和特征信息
            model_info = {
                'model': self.model,
                'feature_columns': self._get_expected_features(),
                'feature_engineer_version': '1.0',
                'training_date': datetime.datetime.now().strftime('%Y-%m-%d')
            }

            joblib.dump(model_info, path)
            print(f"💾 模型已保存: {path}")
            return True
        except Exception as e:
            print(f"❌ 模型保存失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return False

    def load_model(self, model_path=None):
        """加载模型"""

        path = model_path or 'ml/models/score_predictor.pkl'
        try:
            model_info = joblib.load(path)
            self.model = model_info['model']
            print(f"✅ 模型加载成功: {path}")
            print(f"   - 特征数量: {len(model_info['feature_columns'])}")
            print(f"   - 训练日期: {model_info.get('training_date', '未知')}")
            return True
        except Exception as e:
            print(f"❌ 模型加载失败: {str(e)}")
            return False

    def get_model_info(self):
        """获取模型信息"""

        return {
            "status": "loaded" if self.model else "not_trained",
            "model_type": type(self.model).__name__ if self.model else "none",
            "ml_available": True,
            "features": self._get_expected_features(),
            "feature_count": len(self._get_expected_features())
        }

    def evaluate_model(self, test_data=None):
        """评估模型性能"""

        # 这里可以添加模型评估逻辑
        return {
            'status': 'evaluation_ready',
            'message': '模型评估功能就绪'
        }


# 创建训练器实例，方便使用
model_trainer = ModelTrainer()